import bean.Event;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.TimerService;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.PriorityQueue;

/**
 * 时间戳排序函数
 */
public class SortByTimestampFunction extends KeyedProcessFunction<String, Event, Event> {
    private ValueState<PriorityQueue<Event>> queueState = null;

    @Override
    public void open(Configuration config) {
        ValueStateDescriptor<PriorityQueue<Event>> descriptor = new ValueStateDescriptor<>(
                // state name
                "sorted-events",
                // type information of state
                TypeInformation.of(new TypeHint<PriorityQueue<Event>>() {
                }));
        queueState = getRuntimeContext().getState(descriptor);
    }

    @Override
    public void processElement(Event event, Context context, Collector<Event> out) throws Exception {
        TimerService timerService = context.timerService();
        long currentWatermark = timerService.currentWatermark();
//        System.out.format("processElement called with watermark %d\n", currentWatermark);
        if (context.timestamp() > currentWatermark) {
            PriorityQueue<Event> queue = queueState.value();
            if (queue == null) {
                queue = new PriorityQueue<>(10);
            }
            queue.add(event);
            queueState.update(queue);
            timerService.registerEventTimeTimer(event.getTimestamp());
        }
    }

    @Override
    public void onTimer(long timestamp, OnTimerContext context, Collector<Event> out) throws Exception {
        PriorityQueue<Event> queue = queueState.value();
        long watermark = context.timerService().currentWatermark();
//        System.out.format("onTimer called  with watermark %d\n", watermark);
        Event head = queue.peek();
        while (head != null && head.getTimestamp() <= watermark) {
            out.collect(head);
            queue.remove(head);
            head = queue.peek();
        }
    }
}
